package cn.kduck.code.generate.service.impl;

import cn.kduck.code.entity.service.Entity;
import cn.kduck.code.entity.service.EntityField;
import cn.kduck.code.entity.service.EntityService;
import cn.kduck.code.generate.AiProperties;
import cn.kduck.code.generate.CodeGenProgressListener;
import cn.kduck.code.generate.exception.ModelResponseException;
import cn.kduck.code.generate.json.AiChatRequest;
import cn.kduck.code.generate.json.AiChatRequest.ChatMessage;
import cn.kduck.code.generate.json.ChatCompletionChunk;
import cn.kduck.code.generate.json.ChatCompletionChunk.Choice;
import cn.kduck.code.generate.json.ChatCompletionChunk.Delta;
import cn.kduck.code.generate.json.CodeResponse;
import cn.kduck.code.generate.service.GenerateService;
import cn.kduck.code.generate.service.impl.CodeGenProgressListenerImpl.CodeGenProgress;
import cn.kduck.code.module.service.Module;
import cn.kduck.code.module.service.ModuleService;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.commons.compress.archivers.ArchiveOutputStream;
import org.apache.commons.compress.archivers.ArchiveStreamFactory;
import org.apache.commons.compress.archivers.zip.ZipArchiveEntry;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.stereotype.Service;
import org.springframework.util.StringUtils;
import org.springframework.web.client.RestTemplate;

import java.io.*;
import java.util.*;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.function.Consumer;

@Service
public class GenerateServiceImpl implements GenerateService {

    private static final String DATA_PREFIX = "data:";

    private RestTemplate restTemplate = new RestTemplate();
    private ObjectMapper objectMapper = new ObjectMapper();

//    private MessageChatMemory messageChatMemory = new MessageChatMemory();

    private ModuleService moduleService;
    private EntityService entityService;

    private AiProperties aiProperties;

    private CodeGenProgressListener codeGenProgressListener = new CodeGenProgressListenerImpl();

    private ExecutorService executorService = Executors.newFixedThreadPool(5);

    @Value("${kduck.code.gen-base-path}")
    private String genBasePath;

    public GenerateServiceImpl(AiProperties aiProperties,ModuleService moduleService,EntityService entityService){
        this.aiProperties = aiProperties;
        this.moduleService = moduleService;
        this.entityService = entityService;
        objectMapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
    }

    @Override
    public String generate(GenerateType[] types, String moduleId, String prompt) {

        Module module = moduleService.getModule(moduleId);
        executorService.submit(()->{
            String conversationId = "default";
//        List<ChatMessage> messagesList = messageChatMemory.get(conversationId, 20);
//        if(StringUtils.hasText(prompt)){
//            messagesList.add(new ChatMessage("user",prompt));
//        }

            List<Entity> entityList = module.getEntityList();
            StringBuilder entityPromptBuilder = new StringBuilder();
            entityPromptBuilder.append("现有需求：" + module.getRequirement()+"\n");
            entityPromptBuilder.append("生成的代码基础包路径为：" + module.getPackageName() + "。\n");
            entityPromptBuilder.append("生成的Controller基础路径为：" + module.getBaseApiPath() + "。\n");
            entityPromptBuilder.append("现有实体信息：\n");
            for (Entity entity : entityList) {
                entityPromptBuilder.append(entity);

                List<EntityField> entityFieldList = entityService.listEntityField(entity.getEntityId());
                if(!entityFieldList.isEmpty()){
                    entityPromptBuilder.append("\n该实体字段（对应类属性名为表字段名小驼峰形式）：\n");
                    for (EntityField entityField : entityFieldList) {
                        entityPromptBuilder.append(entityField);
                    }
                }
            }

            Map<GenerateType,String[]> codeMap = new HashMap<>();
            for (GenerateType type : types) {
                String finalPrompt = prompt + "。" + entityPromptBuilder;
                if(!codeMap.isEmpty()){
                    StringBuilder stringBuilder = new StringBuilder();
                    for (String[] value : codeMap.values()) {
                        for (String c : value) {
                            stringBuilder.append(c);
                            stringBuilder.append("\r\n---\r\n");
                        }
                    }
                    finalPrompt +="目前已生成的代码有（每个类文件以---分隔）：" + stringBuilder + ",根据目前代码来完善和补充新的生成代码。";
                }
                String basePromptText = loadPromptText(GenerateType.MODULE);
                basePromptText += "一个模块包含若干子包，其中：" + type.getDescribe();

                String systemPromptText = loadPromptText(type);
                List<ChatMessage> messagesList = new ArrayList<>();
                messagesList.add(new ChatMessage("system",basePromptText + "\r\n" +systemPromptText));
                messagesList.add(new ChatMessage("user",finalPrompt));

//        List<ChatMessage> finalChatMessageList = new ArrayList<>(messagesList.size() + 1);
//        finalChatMessageList.add(new ChatMessage("system",basePromptText + "\r\n" +systemPromptText));
//        finalChatMessageList.addAll(messagesList);

                codeGenProgressListener.onProgressUpdate(moduleId,type.name(),null);

                StringBuilder codeBuilder = new StringBuilder();
                sendChatMessage(conversationId, messagesList, response ->{
                    if(response.equals("[DONE]")){
                        CodeResponse[] codeResponse;
                        try {
                            String json = codeBuilder.toString();
                            if(json.startsWith("```json")){
                                json = json.substring(7,json.length() - 3);
                            }
                            codeResponse = objectMapper.readValue(json, CodeResponse[].class);
                            String[] codes = new String[codeResponse.length];
                            for (int i = 0; i < codeResponse.length; i++) {
                                codes[i] = codeResponse[i].getJavaCode();
                            }
                            codeMap.put(type,codes);
                            saveCodeFile(moduleId,type,codeResponse);
                        } catch (JsonProcessingException e) {
                            throw new RuntimeException("语言模型返回的数据格式错误：" + codeBuilder,e);
                        }
                    }else {
                        codeGenProgressListener.onProgressUpdate(moduleId,type.name(),null);
                        codeBuilder.append(response);
                    }
                });
            }
            codeGenProgressListener.onGenerationComplete(moduleId);
        });



        genBasePath = genBasePath.endsWith(File.separator) ? genBasePath: genBasePath + File.separator;
        return genBasePath + module.getPackageName().replaceAll("[.]","\\" + File.separator);
    }

    private void saveCodeFile(String moduleId,GenerateType type,CodeResponse[] codes){
        for (CodeResponse code : codes) {
            String packageInfo = code.getPackageInfo();
            String codePath = StringUtils.replace(packageInfo, ".", File.separator);
            String filePath = codePath + File.separator + code.getClassName() + ".java";
            File codeFile = new File("",codePath);
            if(!codeFile.exists()){
                codeFile.mkdirs();
            }
            File javaCodeFile = new File(genBasePath, moduleId + File.separator + filePath);
            if(javaCodeFile.exists()){
                continue;
            }

            File parentFile = javaCodeFile.getParentFile();
            if(!parentFile.exists()){
                parentFile.mkdirs();
            }

            codeGenProgressListener.onProgressUpdate(moduleId,type.name(),code.getClassName());
            try (FileWriter writer = new FileWriter(javaCodeFile)) {
                writer.write(code.getJavaCode());
            } catch (IOException e) {
                throw new RuntimeException("保存代码文件错误:" + filePath,e);
            }
        }
    }

    private String loadPromptText(GenerateType type){
        StringBuilder contentBuilder = new StringBuilder();
        InputStream promptStream = this.getClass().getResourceAsStream("/cn/kduck/code/generate/prompt/" + type.getPromptName());
        try (BufferedReader reader = new BufferedReader(new InputStreamReader(promptStream))) {
            String line;
            while ((line = reader.readLine()) != null) {
                contentBuilder.append(line).append("\n");
            }
        } catch (IOException e) {
            throw new RuntimeException("读取提示文件错误:" + type.getPromptName(),e);
        }

        return contentBuilder.toString();
    }


    private void sendChatMessage(String conversationId, List<ChatMessage> messagesList, Consumer<String> function) {
        HttpHeaders headers = new HttpHeaders();
        headers.set("Authorization", "Bearer " + aiProperties.getApiKey());
        headers.set("Content-Type", "application/json");

        AiChatRequest request = new AiChatRequest();
//        ChatMessage[] messages = new ChatMessage[]{
//                new ChatMessage("system",SYSTEM_PROMPT),
//                new ChatMessage("user",userPrompt),
//        };
        request.setModel(aiProperties.getModel());
//        request.setToolChoice(aiProperties.getToolChoice().getValue());
        request.setStream(true);
        request.setMessages(messagesList.toArray(new ChatMessage[messagesList.size()]));

        restTemplate.execute(
                aiProperties.getApiUrl(),
                HttpMethod.POST,
                requestCallback -> {
                    requestCallback.getHeaders().addAll(headers);
                    requestCallback.getBody().write(objectMapper.writeValueAsString(request).getBytes());
                },
                responseExtractor -> {
//                    function.accept("[PROCESSING]");
                    InputStream inputStream = responseExtractor.getBody();
                    try (BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream))) {
                        String line;
                        StringBuilder contentBuilder = new StringBuilder();

                        int count = 0;
                        while ((line = reader.readLine()) != null) {
                            if(line.startsWith(DATA_PREFIX)){
                                String dataJson = line.substring(DATA_PREFIX.length()).trim();
                                if(!"[DONE]".equals(dataJson)){
                                    ChatCompletionChunk chatCompletionChunk = objectMapper.readValue(dataJson, ChatCompletionChunk.class);
                                    List<Choice> choices = chatCompletionChunk.getChoices();
                                    if(choices == null){
                                        throw new ModelResponseException(dataJson, "语言模型返回了错误格式的响应：" + dataJson);
                                    }
                                    Optional<Choice> firstChoice = choices.stream().findFirst();
                                    if(firstChoice.isPresent()){
                                        Choice choice = firstChoice.get();
                                        Delta delta = choice.getDelta();
                                        String content = delta.getContent();
                                        if(content.length() > 0){
                                            count++;
                                            System.out.print(count+"\r");
                                            function.accept(content);
                                            contentBuilder.append(content);
                                        }
                                    }
                                }else{
                                        messagesList.add(new ChatMessage("assistant",contentBuilder.toString()));
//                                        messageChatMemory.add(conversationId, messagesList);
                                        function.accept("[DONE]");
                                }
                            }
                        }
                    } catch (IOException e) {
                        throw new RuntimeException("读取请求流错误",e);
                    }

                    return null;
                }
        );
    }

    public void zipCodeFiles(String sourceCodeFilePath,OutputStream outputStream){
        File file = new File(sourceCodeFilePath);
        if(!file.exists()){
            return;
        }
        try (ArchiveOutputStream archive = new ArchiveStreamFactory().createArchiveOutputStream(ArchiveStreamFactory.ZIP, outputStream)) {
            zipCodeFiles(archive, file);
        } catch (Exception e) {
            throw new RuntimeException("打包代码出错",e);
        }
    }

    @Override
    public List<CodeGenProgress> getGenerateProgress(String moduleId) {
        return codeGenProgressListener.getGenProgress(moduleId);
    }

    private void zipCodeFiles(ArchiveOutputStream archive, File file) throws IOException {
        if(file.isDirectory()){
            File[] files = file.listFiles();
            for (File subFile : files) {
                zipCodeFiles(archive,subFile);
            }
        } else {

            String zipPath = StringUtils.replace(file.getPath(),genBasePath,"");
            zipPath = StringUtils.replace(zipPath,File.separator,"/");

            ZipArchiveEntry entry = new ZipArchiveEntry(zipPath);
            archive.putArchiveEntry(entry);
            try (FileInputStream inputStream = new FileInputStream(file)) {
                int len;
                byte[] buffer = new byte[8192];
                while ((len = inputStream.read(buffer)) > -1) {
                    archive.write(buffer, 0, len);
                }
            }
            archive.closeArchiveEntry();
            System.out.println("打包文件：" + zipPath);
        }
    }
}
